import re
import pandas as pd
import json
import time
import datetime
import pytz
from playwright.sync_api import Playwright, sync_playwright, expect
import psycopg2

# 数据库连接配置
DB_CONFIG = {
    'dbname': 'heytea_db',
    'user': 'heytea',  # 确保使用具有足够权限的用户
    'password': '666666',
    'host': '127.0.0.1',
    'port': '5432'
}

def insert_data_to_db(data):
    try:
        # 连接到数据库
        conn = psycopg2.connect(**DB_CONFIG)
        cursor = conn.cursor()

        # 插入数据的SQL语句
        insert_query = """
        INSERT INTO inventory_flow (
            store_id, store_name, code, product_code, product_name, qty, 
            accounting_unit_id, accounting_unit_code, accounting_unit_name, 
            demand_id, demand_qty, demand_hour, spec, order_code, order_type, 
            order_time, action, status,update_time
        ) VALUES (
            %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s,now()
        )
        """

        # 遍历数据并插入到数据库
        for row in data:
            cursor.execute(insert_query, (
                row.get('store_id'), row.get('store_name'), row.get('code'), 
                row.get('product_code'), row.get('product_name'), row.get('qty'), 
                row.get('accounting_unit_id'), row.get('accounting_unit_code'), 
                row.get('accounting_unit_name'), row.get('demand_id'), 
                row.get('demand_qty'), row.get('demand_hour'), row.get('spec'), 
                row.get('order_code'), row.get('order_type'), row.get('order_time'), 
                row.get('action'), row.get('status')
            ))

        # 提交事务
        conn.commit()
        print("数据插入成功")

    except Exception as e:
        print(f"插入数据时发生错误: {e}")

    finally:
        # 关闭数据库连接
        if cursor:
            cursor.close()
        if conn:
            conn.close()

# 定义门店的账号和密码
stores = {
    "东里": {"username": "J21013414", "password": "J@21013414"},
    "莲下": {"username": "J21013524", "password": "J@21013524"}
}

def myPrint(resp) -> None:
    print("<<", resp.status, resp.url, resp.payload)

def run(playwright: Playwright, store_id: str, day_start: str, day_end: str, insert_to_db: bool) -> None:
    browser = playwright.chromium.launch(headless=False)
    context = browser.new_context()
    page = context.new_page()
    page.set_default_timeout(30000)
    page.goto("https://partner-store.heytea.com/")

    store = stores[store_id]
    try:        
        page.get_by_placeholder("请输入账号").fill(store["username"])
        page.get_by_placeholder("请输入密码").fill(store["password"])
        page.get_by_role("button", name="登录").click()
    except:
        print("logined")

    page.get_by_role("link", name="库存流水").first.click()
    page.get_by_placeholder("Start Time").click()
    page.get_by_title(day_start).first.locator("div").click()
    page.get_by_title(day_end).first.locator("div").click()
    page.locator("div").filter(has_text=re.compile(r"^单据类型Select\.\.\.$")).locator(
        "span"
    ).nth(1).click()
    page.get_by_role("option", name="销售").click()

    data = []

    with page.expect_response(re.compile(r".*/inventory/log\?.*")) as response_info:
        page.get_by_role("button", name="查询").click()
    response = response_info.value
    rows = response.json()['payload']['rows']
    data.extend(rows)

    page.get_by_text("条/页").click()
    page.get_by_text("100 条/页").click()
    page.get_by_role("button", name="查询").click()
    

    try:
        while page.get_by_role("button", name="right").count() > 0:
            with page.expect_response(re.compile(r".*/inventory/log\?.*")) as response_info:
                page.get_by_role("button", name="right").click()
            response = response_info.value
            rows = response.json()['payload']['rows']
            data.extend(rows)
    except:
        print("timeouts!")

    for row in data:
        utc_datetime = datetime.datetime.strptime(row['order_time'], "%Y-%m-%dT%H:%M:%SZ")
        local_timezone = pytz.timezone('Asia/Shanghai')
        local_datetime = utc_datetime.replace(tzinfo=pytz.utc).astimezone(local_timezone)
        row['order_time'] = local_datetime.strftime("%Y-%m-%d %H:%M:%S")
        row['order_hour'] = local_datetime.strftime("%H:00:00")
    df = pd.DataFrame(data)
    df.to_excel(f"data_{store_id}_{day_start}_{day_end}.xlsx", index=False)

    # 在生成Excel文件后，根据开关决定是否插入数据到数据库
    if insert_to_db:
        insert_data_to_db(data)

    context.storage_state(path="cway")
    context.close()
    browser.close()

with sync_playwright() as playwright:
    # 获取用户输入的日期范围
    day_start = input("请输入查询起始日期（格式：-MM-DD）：")
    day_end = input("请输入查询结束日期（格式：-MM-DD）：")
    # 获取用户是否要插入数据库的选项
    insert_to_db = input("是否要将数据插入数据库？(y/n): ").lower() == 'y'
    
    # 显示所有可选门店
    print("可选门店：")
    for i, store_id in enumerate(stores.keys(), 1):
        print(f"{i}. {store_id}")
    
    # 获取用户选择的门店
    store_choice = input("请选择要操作的门店编号：")
    store_ids = list(stores.keys())
    selected_store = store_ids[int(store_choice) - 1]
    
    # 只运行用户选择的门店
    run(playwright, selected_store, day_start, day_end, insert_to_db)